from langchain.chains.llm import LLMChain
from langchain_core.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate

from config.model_config import get_chat_openai_xin

template = "您是一位很会将故事的人。"
system_message_prompt = SystemMessagePromptTemplate.from_template(template)
human_template = "讲一个有关{text}的故事"
human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
print(human_message_prompt)


chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])

chain = LLMChain (
    llm = get_chat_openai_xin(streaming=False),
    prompt=chat_prompt,
    # output_parser=CommaSeparatedListOutputParser()
)
# res = chain.run(number = 4, text = "colors")
res = chain.invoke({"text": "老虎"})
print(res)